import numpy as np


def calcInfoGain(feature, label, index):
    '''
    计算信息增益
    :param feature:测试用例中字典里的feature，类型为ndarray
    :param label:测试用例中字典里的label，类型为ndarray
    :param index:测试用例中字典里的index，即feature部分特征列的索引。该索引指的是feature中第几个特征，如index:0表示使用第一个特征来计算信息增益。
    :return:信息增益，类型float
    '''

    #*********** Begin ***********#
    def entropy(labels):
        #这个函数有两个返回值，第一个返回值是数组，数组中只包含唯一值，第二个返回值是也是数组，是唯一值出现的次数
        unique_labels, counts = np.unique(labels, return_counts=True)
        #计算每一个标签出现的概率
        probabilities = counts / len(labels)
        #计算信息熵
        entropy = -np.sum(probabilities * np.log2(probabilities))
        return entropy
    #计算总的信息熵
    total_entropy = entropy(label)
    #求出feature特征数组第index列的唯一值数组和唯一值出现次数的数组
    unique_values, value_counts = np.unique(feature[:, index], return_counts=True)
    weighted_entropy = 0
    #value和count每一次从unique_values和value_counts取一个值进行for循环
    for value, count in zip(unique_values, value_counts):
        #label装的是一个布尔数组，与特征值相等的记为True，不相等的记为False
        subset_label = label[feature[:, index] == value]
        #这里计算的条件熵，注意，计算熵的函数已经带负号了
        weighted_entropy += (count / len(feature)) * entropy(subset_label)
    #计算信息增益
    info_gain = total_entropy - weighted_entropy
    #返回信息增益
    return info_gain

    #*********** End *************#